/**
 * 
 */
package edu.umd.clip.lm.model.data;

import java.io.*;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.nio.channels.FileChannel.MapMode;
import java.util.*;

import edu.umd.clip.lm.util.GrowingOutputBuffer;
import edu.umd.clip.lm.util.Long2IntMap;

public class TrainingDataBlock implements Iterable<ContextFuturesPair> {
	private int size;
	private int bytes;
	private int totalCount;
	private ArrayList<ContextFuturesPair> data;
	private boolean parsed;
	private ByteBuffer buffer;
	
	public static final int DEFAULT_MAX_CAPACITY = 1000000;
	private static int maxCapacity = DEFAULT_MAX_CAPACITY;
	public static final int HEADER_SIZE = 12; 

	public static final TrainingDataBlock TERMINATOR = new TrainingDataBlock();
	
	public TrainingDataBlock() {
		setEmptyHeader();
		parsed = true;
	}
	
	public TrainingDataBlock(int capacity) {
		if (capacity > maxCapacity) capacity = maxCapacity;
		data = new ArrayList<ContextFuturesPair>(capacity);
		parsed = true;
	}
	
	public TrainingDataBlock(DataInput in) throws IOException {
		read(in);
	}
	
	/**
	 * @param input
	 * @throws IOException 
	 */
	public TrainingDataBlock(ReadableByteChannel input) throws IOException {
		ByteBuffer buff = ByteBuffer.allocate(HEADER_SIZE);
		input.read(buff);
		buff.flip();
		size = buff.getInt();
		bytes = buff.getInt();
		totalCount = buff.getInt();
		parsed = true;
		
		//System.err.println("Block read: " + toString());
		
		if (size > 0) {
			// postpone parsing data
			parsed = false;
			data = null;
			/*
			if (bytes > 4*1024*1024) {
				buffer = input.map(MapMode.READ_ONLY, input.position(), bytes);
				input.position(input.position() + bytes);
			} else {
			*/
				buffer = ByteBuffer.allocate(bytes);
				while(buffer.hasRemaining()) {
					input.read(buffer);
				}
				buffer.flip();
			/*
			}
			*/
		}
	}

	void parse() {
		if (parsed) return;
		
		synchronized(this) {
			if (parsed) return;
			
			data = new ArrayList<ContextFuturesPair>(size);
			for(int i=0; i<size; ++i) {
				data.add(new ContextFuturesPair(buffer));
			}		
			buffer = null;
			parsed = true;
		}
	}
	
	public void readHeader(DataInput in) throws IOException {
		size = in.readInt();
		bytes = in.readInt();
		totalCount = in.readInt();
	}
	
	public void writeHeader(DataOutput out) throws IOException {
		out.writeInt(size);
		out.writeInt(bytes);
		out.writeInt(totalCount);
	}
	
	public void writeHeader(GrowingOutputBuffer buffer) {
		buffer.putInt(size);
		buffer.putInt(bytes);
		buffer.putInt(totalCount);
	}
	
	public boolean hasData() {
		return size > 0;
	}
	
	public void setEmptyHeader() {
		size = 0;
		bytes = 0;
		totalCount = 0;
	}
	
	public void read(DataInput in) throws IOException {
		readHeader(in);
		if (hasData()) {
			data = new ArrayList<ContextFuturesPair>(size);
			for(int i=0; i<size; ++i) {
				data.add(new ContextFuturesPair(in));
			}
		}
		parsed = true;
	}
	
	public Long2IntMap makeCounts() {
		parse();
		Long2IntMap map = new Long2IntMap();
		addCounts(map);
		return map;
	}
	
	public void addCounts(Long2IntMap  map) {
		if (!hasData()) return;
		parse();
		for(ContextFuturesPair pair : data) {
			for(TupleCountPair tc : pair.getFutures()) {
				totalCount += tc.count;
				map.addAndGet(tc.tuple, tc.count);
			}
		}
	}
	
	private static class MyByteArrayOutputStream extends ByteArrayOutputStream {
		MyByteArrayOutputStream(int size) {
			super(size);
		}
		// get direct access to avoid copying
		private byte[] getBuf() {
			return this.buf;
		}
	}
	
	public void write(DataOutput out) throws IOException {
		parse();
		if (data == null || data.size() == 0) {
			setEmptyHeader();
			writeHeader(out);
			return;
		}
		
		// TODO: fix this
		TrainingDataBlock.MyByteArrayOutputStream buffer = new MyByteArrayOutputStream(40*size);
		ObjectOutput tmpOut = new ObjectOutputStream(buffer);
		// discard the header written by ObjectOutputStream
		tmpOut.flush();
		buffer.reset();
		
		for(ContextFuturesPair pair : data) {
			pair.writeExternal(tmpOut);
		}
		tmpOut.close();
		
		bytes = buffer.size();
		totalCount = getTotalCount();
		writeHeader(out);
		out.write(buffer.getBuf(), 0, buffer.size());
	}
	
	public void write(GrowingOutputBuffer buffer) {
		parse();
		if (data == null || data.size() == 0) {
			setEmptyHeader();
			writeHeader(buffer);
			return;
		}
		
		GrowingOutputBuffer tmpBuff = new GrowingOutputBuffer(buffer.getBufferSize());
		for(ContextFuturesPair pair : data) {
			pair.write(tmpBuff);
		}
		
		bytes = tmpBuff.size();
		totalCount = getTotalCount();

		//System.err.println("block written: " + toString());
		
		writeHeader(buffer);
		buffer.put(tmpBuff);
	}
	
	public int getTotalCount() {
		if (totalCount == 0) {
			computeTotalCount();
		}
		return totalCount;
	}
	
	protected void computeTotalCount() {
		parse();
		long count = 0;
		for(ContextFuturesPair pair : data) {
			count += pair.getTotalCount();
		}
		this.totalCount = (int) count;
	}

	public boolean add(ContextFuturesPair pair) {
		parsed = true;
		if (data == null) {
			data = new ArrayList<ContextFuturesPair>();
		} else if (data.size() >= maxCapacity) {
			return false;
		}
		boolean result = data.add(pair);
		size = data.size();
		return result;
	}
	/* (non-Javadoc)
	 * @see java.lang.Iterable#iterator()
	 */
	@Override
	public Iterator<ContextFuturesPair> iterator() {
		parse();
		return data == null ? null : data.iterator();
	}

	public static int getMaxCapacity() {
		return maxCapacity;
	}

	public static void setMaxCapacity(int maxCapacity) {
		TrainingDataBlock.maxCapacity = maxCapacity;
	}

	protected ArrayList<ContextFuturesPair> getData() {
		parse();
		return data;
	}
	
	public boolean isTerminator() {
		return size == 0;
	}
	
	public int size() {
		return size;
	}

	public int getByteSize() {
		return bytes;
	}
	
	@Override
	public String toString() {
		return "TrainingDataBlock [bytes=" + bytes + ", size=" + size
				+ ", totalCount=" + totalCount + "]";
	}
}